import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont

classes = ['普通士兵', '工兵', '火箭兵', '指挥官']

# 加载并预处理图片
def load_and_preprocess_image(image_path, target_size=(640, 640)):
    image = Image.open(image_path).convert('RGB')
    original_size = image.size
    image = image.resize(target_size)
    image_array = np.array(image).astype(np.float32) / 255.0
    image_array = np.transpose(image_array, (2, 0, 1))
    image_array = np.expand_dims(image_array, axis=0)
    return image_array, original_size

# 计算两个框的IoU
def calculate_iou(box1, box2):
    x1_min, y1_min, w1, h1 = box1[:4]
    x2_min, y2_min, w2, h2 = box2[:4]
    
    x1_max, y1_max = x1_min + w1, y1_min + h1
    x2_max, y2_max = x2_min + w2, y2_min + h2
    
    inter_x_min = max(x1_min, x2_min)
    inter_y_min = max(y1_min, y2_min)
    inter_x_max = min(x1_max, x2_max)
    inter_y_max = min(y1_max, y2_max)
    
    inter_area = (inter_x_max - inter_x_min) * (inter_y_max - inter_y_min) if inter_x_min < inter_x_max and inter_y_min < inter_y_max else 0
    
    box1_area = w1 * h1
    box2_area = w2 * h2
    
    return inter_area / (box1_area + box2_area - inter_area)

# 合并同类框并保留概率最大的那个
def merge_similar_boxes_by_probability(full_data, threshold=0.9):
    if len(full_data) == 0:
        return []

    merged_data = []
    used = np.zeros(len(full_data), dtype=bool)

    for i in range(len(full_data)):
        if used[i]:
            continue

        current_data = full_data[i]
        current_prob = current_data[-2]
        current_class_idx = int(current_data[-1])

        for j in range(i + 1, len(full_data)):
            if used[j]:
                continue
            
            # 只对同一类别的框进行合并
            if int(full_data[j][-1]) == current_class_idx and calculate_iou(current_data, full_data[j]) > threshold:
                if full_data[j][-2] > current_prob:
                    current_data = full_data[j]
                    current_prob = full_data[j][-2]
                used[j] = True

        merged_data.append(current_data)

    return merged_data

# 进行推理
def run_inference(image_array, onnx_model_path):
    session = ort.InferenceSession(onnx_model_path)
    input_name = session.get_inputs()[0].name
    result = session.run(None, {input_name: image_array})
    return result

# 在图片上绘制矩形框和文本
def draw_boxes_on_image(image, boxes, original_size, font_path, font_size=20, resized_size=(640, 640)):
    draw = ImageDraw.Draw(image)
    width_ratio = original_size[0] / resized_size[0]
    height_ratio = original_size[1] / resized_size[1]
    
    # 加载字体
    font = ImageFont.truetype(font_path, font_size)
    
    for box in boxes:  # Use box instead of using index to reference classes
        x, y, w, h, prob = box[:5]
        x_min = int((x - w / 2) * width_ratio)
        y_min = int((y - h / 2) * height_ratio)
        x_max = int((x + w / 2) * width_ratio)
        y_max = int((y + h / 2) * height_ratio)
        
        draw.rectangle([x_min, y_min, x_max, y_max], outline='red', width=2)
        
        # 中文文本
        class_idx = int(box[-1])  # Use the last element as class index
        text = f"{classes[class_idx]} {prob:.2f}"
        draw.text((x_min, y_min), text, fill='red', font=font)  # 指定字体
        
    return image

# 主程序
def main():
    image_path = r'datasets\images\test\a19.jpg'  # 替换为你的图片路径
    onnx_model_path = 'best.onnx'  # 替换为你的ONNX模型路径
    font_path = './simhei.ttf'  # 替换为你的中文字体路径

    image_array, original_size = load_and_preprocess_image(image_path)
    inference_result = run_inference(image_array, onnx_model_path)

    output_array = inference_result[0]
    if output_array.shape == (1, 8, 8400):
        output_array = np.squeeze(output_array)

        full_data = []
        for idx in range(8400):
            last_four_values = output_array[-4:, idx]
            max_probability = np.max(last_four_values)
            max_index = np.argmax(last_four_values)

            if max_probability > 0.8:
                modified_data = np.concatenate([output_array[:-4, idx], [max_probability, max_index]])
                full_data.append(modified_data)

        merged_data = merge_similar_boxes_by_probability(full_data, threshold=0.8)

        original_image = Image.open(image_path).convert('RGB')

        # 在原图上绘制框
        image_with_boxes = draw_boxes_on_image(original_image, merged_data, original_size, font_path)

        # 显示图像或保存结果
        image_with_boxes.show()  # 或使用 image_with_boxes.save('output_image.jpg')

if __name__ == "__main__":
    main()
